package com.example.demo.util;

import java.awt.image.BufferedImage;
import java.io.File;
import java.util.*;
import java.util.stream.Collectors;
import javax.imageio.ImageIO;

public class PHashCalculator {

    // 相似度阈值（汉明距离≤5视为相似）
    private static final int SIMILARITY_THRESHOLD = 5;

    // 生成pHash指纹
    public static String getPHash(File file) throws Exception {
        BufferedImage img = ImageIO.read(file);

        // 1. 缩放至32x32并灰度化
        BufferedImage resized = resizeAndGray(img, 32, 32);

        // 2. DCT变换
        double[][] dctData = applyDCT(getPixelMatrix(resized));

        // 3. 取8x8低频区域
        double[][] lowFreq = extractLowFrequency(dctData, 8);

        // 4. 计算哈希
        return computeHash(lowFreq);
    }

    // 汉明距离计算
    public static int hammingDistance(String hash1, String hash2) {
        int distance = 0;
        for (int i = 0; i < hash1.length(); i++) {
            if (hash1.charAt(i) != hash2.charAt(i)) {
                distance++;
            }
        }
        return distance;
    }

    // 相似度转换（0~1）
    public static double similarity(String hash1, String hash2) {
        int maxLength = Math.max(hash1.length(), hash2.length());
        return 1 - (hammingDistance(hash1, hash2) * 1.0 / maxLength);
    }

    //--- 私有辅助方法 ---//
    private static BufferedImage resizeAndGray(BufferedImage src, int w, int h) {
        BufferedImage dest = new BufferedImage(w, h, BufferedImage.TYPE_BYTE_GRAY);
        dest.getGraphics().drawImage(src.getScaledInstance(w, h, java.awt.Image.SCALE_SMOOTH), 0, 0, null);
        return dest;
    }

    private static double[][] getPixelMatrix(BufferedImage img) {
        int width = img.getWidth();
        int height = img.getHeight();
        double[][] matrix = new double[width][height];
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                matrix[x][y] = (img.getRGB(x, y) & 0xFF) / 255.0; // 获取灰度值
            }
        }
        return matrix;
    }

    // DCT变换实现
    private static double[][] applyDCT(double[][] matrix) {
        int N = matrix.length;
        double[][] dct = new double[N][N];
        double alpha = Math.sqrt(2.0/N);

        for (int u = 0; u < N; u++) {
            for (int v = 0; v < N; v++) {
                double sum = 0.0;
                double cu = 0;
                double cv = 0;
                for (int i = 0; i < N; i++) {
                    for (int j = 0; j < N; j++) {
                        cu = (u == 0) ? 1 / Math.sqrt(2) : 1;
                        cv = (v == 0) ? 1 / Math.sqrt(2) : 1;
                        sum += matrix[i][j] * Math.cos((2 * i + 1) * u * Math.PI / (2 * N))

                                * Math.cos((2 * j + 1) * v * Math.PI / (2 * N));
                    }
                }
                dct[u][v] = alpha * cu * cv * sum;
            }
        }
        return dct;
    }

    private static double[][] extractLowFrequency(double[][] dct, int size) {
        double[][] low = new double[size][size];
        for (int i = 0; i < size; i++) {
            System.arraycopy(dct[i], 0, low[i], 0, size);
        }
        return low;
    }

    private static String computeHash(double[][] matrix) {
        double total = 0;
        for (double[] row : matrix) {
            for (double val : row) {
                total += val;
            }
        }
        double avg = total / (matrix.length * matrix[0].length);

        StringBuilder hash = new StringBuilder();
        for (double[] row : matrix) {
            for (double val : row) {
                hash.append(val > avg ? "1" : "0");
            }
        }
        return hash.toString();
    }

    // 核心搜索方法
    public static List<SearchResult> searchSimilarImages(File targetFile, File folder) {
        try {
            String targetHash = PHashCalculator.getPHash(targetFile);
            return Arrays.stream(Objects.requireNonNull(folder.listFiles()))
                    .parallel()  // 并行处理提升性能
                    .filter(f -> f.isFile() && isImageFile(f))
                    .map(f -> {
                        try {
                            String hash = PHashCalculator.getPHash(f);
                            int distance = PHashCalculator.hammingDistance(targetHash, hash);
                            double similarity = PHashCalculator.similarity(targetHash, hash);
                            return new SearchResult(f, distance, similarity);
                        } catch (Exception e) {
                            return null;  // 跳过无法处理的文件
                        }
                    })
                    .filter(Objects::nonNull)
                    .filter(r -> r.hammingDistance <= SIMILARITY_THRESHOLD)
                    .sorted(Comparator.comparingInt(r -> r.hammingDistance))
                    .collect(Collectors.toList());
        } catch (Exception e) {
            throw new RuntimeException("图片处理失败", e);
        }
    }

    // 判断是否为图片文件
    private static boolean isImageFile(File file) {
        String[] imgExtensions = {"jpg", "jpeg", "png", "bmp", "gif"};
        return Arrays.stream(imgExtensions)
                .anyMatch(ext -> file.getName().toLowerCase().endsWith("." + ext));
    }

    // 结果存储类
    static class SearchResult {
        File file;
        int hammingDistance;
        double similarity;

        public SearchResult(File file, int hd, double sim) {
            this.file = file;
            this.hammingDistance = hd;
            this.similarity = sim;
        }
    }

    public static void main(String[] args) throws Exception {
        File img1 = new File("D:/input/01-05.jpg");
        File img2 = new File("D:/input1/01-01.jpg");

        String hash1 = PHashCalculator.getPHash(img1);
        String hash2 = PHashCalculator.getPHash(img2);

        int distance = PHashCalculator.hammingDistance(hash1, hash2);
        double similarity = PHashCalculator.similarity(hash1, hash2);

        System.out.println("汉明距离: " + distance);
        System.out.println("相似度: " + String.format("%.2f%%", similarity * 100));
    }
}